#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <omp.h>
#include "util.h"

int main(int argc, char *argv[]) {
  int my_rank;
  int i,j, n_jobs;
  int nameSize;
  int num_threads;
  char node[MPI_MAX_PROCESSOR_NAME];

  double start_time;
  
  /* Variáveis usadas no controle do loop */
  int continua = 1;
  int iter = 0;

  int lines_per_slave; /* Armazena quantos X[i] cada escravo vai calcular */
  int lines_left;      /* Armazena quantos X[i] o nó mestre vai calcular */

  double *X, *Xk;      /* Vetores de iteração */
  double **A, *b;      /* Matrizes */

  MPI_Status status;   /* Usado nas funções de send e receive */

  /* Estrutura que mantém toda informação necessária para a resolução do sistema */
  info *sistema;

  /* Inicialização MPI */
  MPI_Init(&argc,&argv);
  
  MPI_Comm_rank(MPI_COMM_WORLD,&my_rank); /* Pega o rank do processo */
  MPI_Comm_size(MPI_COMM_WORLD,&n_jobs);  /* Pega o número total de processos */
  MPI_Get_processor_name(node,&nameSize); /* Nome do hospedeiro */

  if(argc > 2) {
    sscanf(argv[2],"%d",&num_threads);
  } else {
    exit(EXIT_FAILURE);
  }

  /* Parsing do arquivo do sistema */
  if(my_rank == 0) { /* Se for o mestre, faz o parsing */
    if(argc > 1) {      
      sistema = createInfo(argv[1]);     
      /* Se aqui forem enviados apenas os seeks que cada nó irá precisar, em vez da tabela inteira, tempo seria economizado */
      broadcastInfo(sistema);
    }
  } else { /* Se for escravo, recebe o parsing do mestre */
    if(argc > 1) {
      sistema = receiveInfo(argv[1]);
    }
  }

  /* X é o único vetor que é mantido inteiro nos processos */
  X = (double *)malloc(sistema->order*sizeof(double));
  for(i=0;i<sistema->order;i++)
    X[i] = 0.0;

  /* Cálculo para distribuição de trabalho */
  lines_per_slave = sistema->order / (n_jobs - 1);
  lines_left = sistema->order % (n_jobs - 1);

  printf("Process %d started at %s\n",my_rank,node);

  if(my_rank == 0) { /* Master */
    int *lines;
    int start_index;

    /* Alocação do vetor Xk no mestre. Nos escravos ele não é alocado inteiro */
    Xk = (double *)malloc(sistema->order*sizeof(double));

    printf("--------------------\nJobs: %d\nThreads per jobs: %d\nLines per slave: %d\nLines calculated here: %d\n-------------------\n",n_jobs,num_threads,lines_per_slave,sistema->order % (n_jobs - 1));
    
    /* Distribuição de trabalho para os nós escravos */
    for(i=1;i<n_jobs;i++) {
      int *lines = (int *)malloc(lines_per_slave*sizeof(int)); /* Esse vetor contém os índices das váriaveis a serem calculadas */
      for(j=0;j<lines_per_slave;j++) {        
        int index = (i-1)*lines_per_slave + j; /* O indíce da j-ésima variável calculada pelo nó i */
        lines[j] = index;
      }
      MPI_Send(lines,lines_per_slave,MPI_INT,i,0,MPI_COMM_WORLD); /* Envia o array inteiro */
    }

    /*
      Se a divisão entre o número de variáveis e o número de escravos não é exata
      o resto das variáveis são calculadas no nó mestre
    */
    if(lines_left != 0) {
      /* Alocação de recursos */
      lines = (int *)malloc(lines_left*sizeof(int));
      b = (double *)malloc(lines_left*sizeof(double));
      A = (double **)malloc(lines_left*sizeof(double *));    

      /* Armazena o primeiro índice a ser calculado no mestre */
      start_index = (i-1)*lines_per_slave;

      /* O array lines aqui funciona exatamente igual ao presente nos nós escravos */
      for(j=0;j<lines_left;j++) {
        int index = (i-1)*lines_per_slave + j;
        lines[j] = index;
      }
      
      /* Carrega apenas as linhas necessárias da matriz A */
      printf("Carregando linhas [%d ... %d] de A em %d\n",lines[0],lines[lines_left-1],my_rank);
      for(i=0;i<lines_left;i++) {
        A[i] = getLine(sistema,lines[i]);
      }
      printf("Terminou de carregar linhas [%d ... %d] de A em %d\n",lines[0],lines[lines_left-1],my_rank);

      /* Carrega apenas as linhas necessárias da matriz b */
      for(i=0;i<lines_left;i++) {
        b[i] = getValue(sistema,-1,lines[i]);
      }
    }

    printf("Starting...\n");
    start_time = omp_get_wtime();
    do {
      /* Primeira coisa: transmitir o vetor X para todos os escravos */        
      MPI_Bcast(X,sistema->order,MPI_DOUBLE,my_rank,MPI_COMM_WORLD);

      if(lines_left != 0) {
        /*
          Calcula o resto das váriaveis no próprio mestre.
          Note que o vetor Xk usado aqui é o Xk global, do próprio mestre, portanto
          a variável start_index precisa ser usada para o cálculo das posições corretas
          de armazenamento dos resultados
        */
        for(i=0;i<lines_left;i++) {
          Xk[i+start_index] = 0.0;
          for(j=0;j<sistema->order;j++) {
            if(lines[i] != j) {
              Xk[i+start_index] += A[i][j]*X[j];
            }
          }

          Xk[i+start_index] = -Xk[i+start_index];
          Xk[i+start_index] += b[i];
          Xk[i+start_index] = Xk[i+start_index]/A[i][lines[i]];
        }

      }
      
      /* Recebe e armazena no vetor Xk os resultados provenientes do trabalho dos escravos */
      for(i=1;i<n_jobs;i++) {
        MPI_Recv(&Xk[(i-1)*lines_per_slave],lines_per_slave,MPI_DOUBLE,i,0,MPI_COMM_WORLD,&status);
      }

      iter++;
      if(iter%100 == 0)
        fprintf(stderr,"Iter: %d\n",iter);

      /* Cálculo do erro e verificação do número de iterações */
      if(iter >= sistema->MAX_ITER || parada(sistema->erro,X,Xk,sistema->order))
        continua = 0;

      /* Manda a mensagem de controle "continua" por broadcast */
      MPI_Bcast(&continua,1,MPI_INT,my_rank,MPI_COMM_WORLD);
            
      swap(X,Xk,sistema->order);
    }
    while (continua);

    printf("Time calculating: %.10lf\n",(omp_get_wtime() - start_time)/60.0);

    printf("Finished in %d iterations\n",iter);
    printf("Test on the row %d: %.5lf =? %.5lf\n",sistema->test_row,calculateLineDistributed(sistema,X),getValue(sistema,-1,sistema->test_row));

  } else { /* Todos os escravos */
    int *lines = (int *)malloc(lines_per_slave*sizeof(int));   
    int t_id;

    b = (double *)malloc(lines_per_slave*sizeof(double));
    A = (double **)malloc(lines_per_slave*sizeof(double *));    
    Xk = (double *)malloc(lines_per_slave*sizeof(double));

    /**
      Recebe os índices da linha
    */
    MPI_Recv(lines,lines_per_slave,MPI_INT,0,0,MPI_COMM_WORLD,&status);

    /**
      Carrega somente as linhas necessarias da matriz A para o escravo atual
    */
    printf("Carregando linhas [%d ... %d] de A em %d\n",lines[0],lines[lines_per_slave-1],my_rank);
    for(i=0;i<lines_per_slave;i++) {
      A[i] = getLine(sistema,lines[i]);
    }
    printf("Terminou de carregar linhas [%d ... %d] de A em %d\n",lines[0],lines[lines_per_slave-1],my_rank);

    /**
      Carrega somente as linhas necessarias da matriz B para o escravo atual
    */
    for(i=0;i<lines_per_slave;i++) {
      b[i] = getValue(sistema,-1,lines[i]);
    }

    omp_set_num_threads(num_threads);
    /* Definição da região paralela. Definindo ela aqui evita-se o overhead de criação de threads em toda iteração.
       Dessa forma as threads são criadas apenas uma vez e permanessem vivas durante toda execução do processo */
    #pragma omp parallel shared(lines_per_slave,sistema,lines,Xk,A,X,b,continua) private(i,j,t_id)
    {
      do {        
        int lines_per_thread = lines_per_slave / num_threads;
        int limit = lines_per_thread;
        t_id = omp_get_thread_num();

        /* Apenas uma thread (a que chegar primeiro aqui) executa a função de recebimento do vetor
           da iteração anterior */
        #pragma omp single
        {
          MPI_Bcast(X,sistema->order,MPI_DOUBLE,0,MPI_COMM_WORLD);
        }      
        /* Aqui há uma barreira natural imposta pela especificação da openMP, ou seja,
           não há perigo de uma thread passar daqui antes de o vetor X estar carregar no nó */

        /* Define quantas variáveis serão calculadas por thread */
        if(t_id == num_threads - 1) {
          limit += lines_per_slave % num_threads;
        }

        /* Cada thread calcula sua parte do trabalho */
        for(i=t_id*lines_per_thread;i<t_id*lines_per_thread+limit;i++) {
          Xk[i] = 0.0;
          for(j=0;j<sistema->order;j++) {
            if(lines[i] != j) {
              Xk[i] += A[i][j]*X[j];
            }
          }

          Xk[i] = -Xk[i];
          Xk[i] += b[i];
          Xk[i] = Xk[i]/A[i][lines[i]];
        }            
     
        /* Barreira para que o vetor Xk só seja enviado depois que ele estiver completamente calculado */
        #pragma omp barrier
     
        /* Apenas uma thread envia o vetor Xk, e recebe o valor de continua.
           continua é uma variável compartilhada por todas as threads, portanto
           quando chegar a sinal de parar, todas as threads terminam. */
        #pragma omp single
        {
          /* Envia seu resultado para o nó mestre e aguarda pela mensagem de controle */
          MPI_Send(Xk,lines_per_slave,MPI_DOUBLE,0,0,MPI_COMM_WORLD);      
          MPI_Bcast(&continua,1,MPI_INT,0,MPI_COMM_WORLD);
        }

      } while(continua);
    }
  }

  MPI_Finalize();

  return EXIT_SUCCESS;
}
